/*
 *    Copyright 2009 Vanessa Williams
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *        http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */

/*
 * KMerFrequency.java
 * Created on Mar 13, 2010
 */

package bioformatica.annotation;

import bioformatica.sequence.FASTASequence;
import bioformatica.util.CloseableIterator;
import org.apache.commons.cli.*;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.*;
import java.util.logging.Logger;

/**
 * Main class to calucluate k-mer frequencies and odds-ratios (to aid in the discovery of unusual k-mers.)
 *
 * Takes the following options on the command line:
 * <ul>
 * <li>-k the number of bases to consider sequentially, e.g. k=2 to study dimers/li>
 * <lu>=seq the FAST sequence file to analyse
 * </ul>
 */
public class KMerFrequency {

    private static final Logger LOG = Logger.getLogger(KMerFrequency.class.getName());

    private int baseCount = 0;
    private int k = 0;
    private String fastaFilename;
    private FASTASequence sequence;


    /**
     * Constructor
     * @param fastaFilename the name of the FASTA file containing the sequence
     * @param k the k in k-mer
     */
    public KMerFrequency(String fastaFilename, Integer k) {

        try {
            this.fastaFilename = fastaFilename;
            this.sequence = new FASTASequence(new File(fastaFilename));
            this.k = k;
        } catch (FileNotFoundException e) {
            LOG.warning("File not found: " + fastaFilename);
        } catch (IOException e1) {
            LOG.warning("Unable to create file: " + fastaFilename);
        }

    }

    /**
     * Calculate the counts for every k-mer appearing in the sequence.
     * Only non-ambiguous bases (A,C,G,T) are considered.
     * @return a Map of Strings (k-mers) and Integer counts
     */
    public Map<String, Integer> kMerCounts() {

        Map<String, Integer> kmerCounts = new TreeMap<String, Integer>();

        CloseableIterator<Byte> bytes = null;
        try {
            bytes = sequence.iterator();
            byte b;
            byte[] bases = new byte[k];
            if (bytes.hasNext()) b = bytes.next();
            baseCount++;
            while(bytes.hasNext()) {
                for (int index=1; index<k; index++) {
                    if (bytes.hasNext()) {
                        b = bytes.next();
                        bases[index]=b;
                        baseCount++;
                    }
                }
                if (bases.length < k) {
                    // we've reached the end of the sequence
                    // and can't do anything with the short sequence
                    // of bases ( < k)
                    break;
                }

                // make a string from our  retrieved bytes
                String kmer = new String(bases);

                // exclude kmers with ambiguous bases
                boolean kmerGood = true;
                char[] kmerchars = kmer.toCharArray();
                for (int i=0; i<kmerchars.length; i++) {
                    char c = kmerchars[i];
                    if (!(c == 'A' || c == 'C' || c == 'G' || c =='T')) {
                        kmerGood = false;
                        break;
                    }
                }
                if (kmerGood) {
                    Integer count = kmerCounts.get(kmer);
                    if (count != null) {
                        kmerCounts.put(kmer, count+1);
                    }
                    else {
                        kmerCounts.put(kmer, 1);
                    }
                }

                bases[0] = bases[1];
            }
        }
        finally {
            if (bytes != null) bytes.close();
        }

        return kmerCounts;
    }

    /**
     * A method to print a report of the calculated counts. For dimers
     * it will print a matrix of frequencies.
     *
     * @param kmerCounts a Map of k-mers and counts (output from kMerCounts)
     * @return a String suitable for printing
     */
    public String kMerMapToString(Map<String, Integer> kmerCounts) {
        StringBuilder sb = new StringBuilder();
        sb.append("kmers = \n\n");
        Set<String> kmers = kmerCounts.keySet();
        for (String key : kmers) {
            sb.append("\t" + key + " : " + kmerCounts.get(key) + "\n");
        }
        sb.append("\n");
        if (k == 2) {
            // output a matrix
            sb.append("matrix = \n\n");
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("AA")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("AC")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("AG")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("AT")/baseCount) + "\n");

            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("CA")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("CC")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("CG")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("CT")/baseCount) + "\n");

            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("GA")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("GC")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("GG")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("GT")/baseCount) + "\n");

            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("TA")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("TC")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("TG")/baseCount));
            sb.append("\t" + String.format("%.4f", (float)kmerCounts.get("TT")/baseCount) + "\n");

        }
        return sb.toString();
    }

    /**
     * Parses the command-line arguments, calls kMerCounts() and kMerMapToString()
     * and writes the result to standard output.
     *
     * @param args command-line arguments
     */
    public static void main (String[] args) {
        Options options = new Options();
        Option command = OptionBuilder.withArgName("n").
                                        hasArg().
                                        withDescription("Specify the number of contiguous bases to test frequencies for").
                                        create("k");
        Option file = OptionBuilder.withArgName("file").
                                        hasArg().
                                        withDescription("FASTA file name").
                                        create("seq");
        options.addOption(command);
        options.addOption(file);

        String k="null";
        CommandLineParser parser = new GnuParser();
        try {
            // parse the command line arguments
            CommandLine line = parser.parse( options, args );
            String output = null;
            if (line.hasOption("k")) {
                k = line.getOptionValue("k");
            }
            if (line.hasOption("seq")) {
                String fasta = line.getOptionValue("seq");
                KMerFrequency instance = new KMerFrequency(fasta, Integer.parseInt(k));
                Map<String, Integer> kmercounts = instance.kMerCounts();
                output = instance.kMerMapToString(kmercounts);
                System.out.println(output);
            }
        }
        catch( ParseException exp ) {
            // oops, something went wrong
            LOG.exiting( "Parsing failed.  Reason: ", exp.getMessage() );
            System.exit(1);
        }

    }

}
